package com.chaos.framework.model.helper;

import cn.hutool.core.util.IdUtil;
import com.chaos.framework.model.constant.Constants;
import org.slf4j.MDC;

import java.util.Objects;
import java.util.concurrent.Callable;


public class TraceHelper {

    public static String getTraceId() {
        String traceId = MDC.get(Constants.TRACE_ID);

        if (Objects.isNull(traceId)) {
            traceId = IdUtil.simpleUUID();
        }
        return traceId;
    }

    public static void setTraceId() {
        MDC.put(Constants.TRACE_ID, IdUtil.simpleUUID());
    }

    public static void clearTraceId() {
        MDC.clear();
    }

    public static Runnable wrapper(Runnable r) {
        return () -> {
            setTraceId();
            try {
                r.run();
            } finally {
                clearTraceId();
            }
        };
    }

    public static <T> Callable<T> wrapper(Callable<T> c) {
        return () -> {
            setTraceId();
            try {
                return c.call();
            } finally {
                clearTraceId();
            }
        };
    }
}
